import numpy as np
import math

def pred(X, P, dt):
    noise_ax = 9
    noise_ay = 9
    F = np.array([[1, 0, dt, 0],
                  [0, 1, 0, dt],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]], dtype=np.float)
    Q = np.array([[dt**4/4*noise_ax, 0, dt**3/2*noise_ax, 0],
                  [0, dt**4/4*noise_ay, 0, dt**3/2*noise_ay],
                  [dt**3/2*noise_ax, 0, dt**2*noise_ax, 0],
                  [0, dt**3/2*noise_ay, 0, dt**2*noise_ay]], dtype=np.float)
    X_pred = np.matmul(F, X)
    P_pred = np.matmul(np.matmul(F, P), F.transpose()) + Q
    return X_pred, P_pred

def laser_update(X, P, Z):
    H = np.array([[1, 0, 0, 0],
                  [0, 1, 0, 0]], dtype=np.float)
    R = np.array([[0.0225, 0],
                  [0, 0.0225]], dtype=np.float)
    Ht = np.transpose(H)
    K = np.matmul(np.matmul(P,Ht), np.linalg.inv(np.matmul(np.matmul(H, P), Ht)+R))
    
    I = np.identity(4) 
    X_update = X + np.matmul(K,(Z - np.matmul(H,X)))
    P_update = np.matmul((I-np.matmul(K,H)),P)
    return X_update, P_update

def radar_update(X, P, Z):
    H = jacobian(X)
    R = np.array([[0.09, 0, 0],
                  [0, 0.0009, 0],
                  [0, 0, 0.09]], dtype=np.float)
    Ht = np.transpose(H)
    K = np.matmul(np.matmul(P,Ht), np.linalg.inv(np.matmul(np.matmul(H, P), Ht)+R))
    
    I = np.identity(4) 
    X_update = X + np.matmul(K,(Z - cartesian2polar(X)))
    P_update = np.matmul((I-np.matmul(K,H)),P)
    return X_update, P_update

def cartesian2polar(X):
    px = X[0][0]
    py = X[1][0]
    vx = X[2][0]
    vy = X[3][0]
    rho = math.sqrt(px*px + py*py);
    phi = math.atan(py/px);
    if px<0 and py>0:
        phi += math.pi
    elif px<0 and py<0:
        phi -= math.pi
        
    if rho < 0.000001:
        rho = 0.000001
    rho_dot = (px*vx + py*vy) / rho
    X_Polar = np.array([[rho], [phi], [rho_dot]], dtype=np.float)
    return X_Polar

def jacobian(X):
    px = X[0][0]
    py = X[1][0]
    vx = X[2][0]
    vy = X[3][0]
  
    c1 = px*px+py*py
    c2 = math.sqrt(c1)
    c3 = c1*c2
    Hj = np.zeros((3,4))
    if abs(c1) < 0.0001:
        print("CalculateJacobian () - Error - Division by Zero")
        
    else:
        Hj[0][0] = px/c2
        Hj[0][1] = py/c2
        Hj[1][0] = -py/c1
        Hj[1][1] = px/c1
        Hj[2][0] = py*(vx*py - vy*px)/c3
        Hj[2][1] = px*(px*vy - py*vx)/c3
        Hj[2][2] = px/c2
        Hj[2][3] = py/c2

    return Hj























